package com.jeremy.comutils.utils;

import java.io.*;
import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.UnknownHostException;
import java.text.NumberFormat;

/***
 * Tool to measure TCP throughput, similar to iperf
 * @author Bela Ban
 * @version $Id: JPerf.java,v 1.4 2007/12/28 23:29:17 belaban Exp $
 */
public class Jperf implements Runnable {
    boolean client;
    boolean direction;//test direction:true-up test,false-down
    int num;
    InetAddress local_addr;
    int local_port;
    int remote_port;
    InetAddress remote_addr;
    int size;
    int receivebuf = 200000;
    int sendbuf = 200000;
    static NumberFormat f;
    Socket client_sock;

    static {
        f = NumberFormat.getNumberInstance();
        f.setGroupingUsed(false);
        f.setMaximumFractionDigits(2);
    }

    public Jperf() {
    }

    //实现Thread类的接口Runnable， 用于支持服务端连接的多线程
    @Override
    public void run() {

        try {
            server_accept_data();
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }

    }

    //用于applet和GUI调用
    public Jperf(String remote_addr, int remote_port, int num, int size, boolean direction) throws UnknownHostException {

        this.remote_addr = InetAddress.getByName(remote_addr);
        this.remote_port = remote_port;
        this.num = num;
        this.size = size;
        this.direction = direction;
    }

    private void start(boolean client, boolean direction, int num, int size, String local_addr, int local_port,
                       String remote_addr, int remote_port, int receivebuf, int sendbuf) throws IOException, ClassNotFoundException {
        this.client = client;
        this.direction = direction;
        this.num = num;
        this.size = size;
        this.local_addr = InetAddress.getByName(local_addr);
        this.local_port = local_port;
        this.remote_addr = InetAddress.getByName(remote_addr);
        this.remote_port = remote_port;
        this.receivebuf = receivebuf;
        this.sendbuf = sendbuf;

        if (client) {
            client();
        } else {
            server();
        }
    }

    //客户端调用
    public String client() throws IOException {

        System.out.println("-- creating socket to " + this.remote_addr + ":" + this.remote_port);
        client_sock = new Socket(this.remote_addr, remote_port);

        String result;
        if (sendExchangeData()) {
            if (direction) result = sendData(num, size);
            else result = receiveData(num, size);
        } else {
            result = "connect to server and exchange data fail!";
        }
        System.out.println(result);
        client_sock.close();

        return result;
    }

    //客户端向服务端发送测试 的数据参数
    private boolean sendExchangeData() throws IOException {
        boolean ret = true;

        client_sock.setReceiveBufferSize(receivebuf);
        client_sock.setSendBufferSize(sendbuf);
        ObjectOutputStream write = new ObjectOutputStream(new BufferedOutputStream(client_sock.getOutputStream()));
        write.writeObject(num);
        write.flush();
        write.writeObject(size);
        write.flush();
        write.writeObject(direction);
        write.flush();

        return ret;
    }

    //服务端调用
    private void server() throws IOException, ClassNotFoundException {
        ServerSocket srv_sock = new ServerSocket(local_port, 10, this.local_addr);
        System.out.println("-- waiting for client on " + srv_sock.getLocalSocketAddress());
        while (true) {
            //wait for a client connect
            Socket client_sock = srv_sock.accept();
            //start a new thread deal this connection:
            Jperf thread_client = new Jperf();
            thread_client.client_sock = client_sock;
            thread_client.sendbuf = sendbuf;
            thread_client.receivebuf = receivebuf;
            //每一个客户端单独一个线程，支持多个客户端同时连接
            Thread thread = new Thread(thread_client);
            thread.start();
        }

    }

    //服务器接收和发送测试 数据
    private void server_accept_data() throws IOException, ClassNotFoundException {
        client_sock.setReceiveBufferSize(receivebuf);
        client_sock.setSendBufferSize(sendbuf);
        System.out.println("-- accepted data connection from " + client_sock.getRemoteSocketAddress());
        ObjectInputStream in = new ObjectInputStream(new BufferedInputStream(client_sock.getInputStream()));

        int num, size;
        boolean direction;
        num = (Integer) in.readObject();
        size = (Integer) in.readObject();
        direction = (Boolean) in.readObject();

        if (num > 0 && size > 0) {
            String result;
            if (direction) result = receiveData(num, size);
            else result = sendData(num, size);

            System.out.println(result);
        } else {
            System.out.println("-- invalid exchange data! ");
        }
        client_sock.close();
    }

    //发送数据，并计算测试结果
    private String sendData(int num, int size) throws IOException {
        System.out.println("-- sending data to " + client_sock.getRemoteSocketAddress().toString() + " total " + num + " messages");
        DataOutputStream out = new DataOutputStream(new BufferedOutputStream(client_sock.getOutputStream()));
        byte[] buf = new byte[size];
        for (int i = 0; i < buf.length; i++) buf[i] = (byte) (i % 128);
        long start, stop;
        int cnt = 1;
        int incr = num / 10;
        start = System.currentTimeMillis();
        for (int i = 0; i < num; i++) {
            out.write(buf, 0, buf.length);
            out.flush();
            if (cnt % incr == 0)
                System.out.println("-- sent " + cnt + " messages");
            cnt++;
        }
        stop = System.currentTimeMillis();
        long diff = stop - start;

        return report("发送报文至 " + client_sock.getRemoteSocketAddress().toString(), (long) num * (long) size, diff);
    }

    //接收数据，并计算测试结果
    private String receiveData(int num, int size) throws IOException {
        System.out.println("-- accepted data from " + client_sock.getRemoteSocketAddress().toString() + " total " + num + " messages");
        DataInputStream in = new DataInputStream(new BufferedInputStream(client_sock.getInputStream()));
        byte[] buf = new byte[size];
        long counter = 0;
        int incr = num / 10;
        long start = 0, stop;
        while (true) {
            int len = in.read(buf, 0, buf.length);
            if (len <= 0) break;
            if (start == 0)
                start = System.currentTimeMillis();
            counter += len;
            if ((counter / size) % incr == 0)
                System.out.println("-- received " + counter / size);
        }
        stop = System.currentTimeMillis();
        long diff = stop - start;

        return report("接收报文来自 " + client_sock.getRemoteSocketAddress().toString(), counter, diff);
    }

    //计算测试结果
    private String report(String direction, long totalbyte, double diff) {
        StringBuilder sb = new StringBuilder();
        double tbs = totalbyte / (1024 * 1024);
        if (tbs < 1000)
            sb.append("\n").append(direction).append("，测试数据总数").append(f.format(tbs)).append("Mbyte").append(" ，用时 ").append(diff).append("毫秒 ");
        else {
            tbs = tbs / 1024;
            sb.append("\n").append(direction).append("，测试数据总数").append(f.format(tbs)).append("Gbyte").append(" ，用时 ").append(diff).append("毫秒 ");
        }
        //tcp throughput:
        double throughput = totalbyte / (diff / 1000.0) / 1024.0;
        if (throughput < 1000)
            sb.append("\n网络吞吐量: ").append(f.format(throughput)).append("KB/秒");
        else {
            throughput /= 1024.0;
            sb.append("\n网络吞吐量: ").append(f.format(throughput)).append("MB/秒");
        }
        //bandwidth
        double bandwidth = totalbyte / (diff / 1000.0) / 1024.0 * 8;
        if (bandwidth < 1000) {
            sb.append("\n网络带宽: ").append(f.format(bandwidth)).append("Kb/秒");
        } else {
            bandwidth /= 1024.0;
            if (bandwidth > 1000) {
                bandwidth /= 1024;
                sb.append("\n网络带宽: ").append(f.format(bandwidth)).append("Gb/秒");
            } else sb.append("\n网络带宽: ").append(f.format(bandwidth)).append("Mb/秒");
        }
        return sb.toString();
    }

    static void help() {
        System.out.println("JPerf [-help] [-client] [-direction <up|down>] [-num <number of msgs] [-size <bytes>] [-local_addr <interface>] [-local_port <port]" +
                "[-remote_addr <IP addr>] [-remote_port <port>] [-receivebuf <bytes>] [-sendbuf <bytes>]");
    }

    //主程序执行入口
    public static void main(String[] args) throws UnknownHostException {
        boolean client = false;
        boolean direction = false;//test direction:true-up test,false-down test
        int num = 10000;
        int size = 8192;
        String local_addr = InetAddress.getLocalHost().getHostAddress();
        String remote_addr = local_addr;
        int local_port = 5000;
        int remote_port = 5000;
        int receivebuf = 200000, sendbuf = 200000;
        for (int i = 0; i < args.length; i++) {
            if (args[i].equals("-client")) {
                client = true;
                continue;
            }
            if (args[i].equals("-num")) {
                num = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-size")) {
                size = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-local_addr")) {
                local_addr = args[++i];
                continue;
            }
            if (args[i].equals("-remote_addr")) {
                remote_addr = args[++i];
                continue;
            }
            if (args[i].equals("-local_port")) {
                local_port = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-remote_port")) {
                remote_port = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-receivebuf")) {
                receivebuf = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-sendbuf")) {
                sendbuf = Integer.parseInt(args[++i]);
                continue;
            }
            if (args[i].equals("-direction")) {
                String value = args[++i];
                direction = value.toLowerCase().equals("up");

                continue;
            }
            help();
            return;
        }
        try {
            new Jperf().start(client, direction, num, size, local_addr, local_port, remote_addr, remote_port, receivebuf, sendbuf);
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
    }
}